from tqdm import tqdm

from pasers.paser_utils import traverse_dir, re_filter
import json
import os


# from preprocess.seal_recog import seal_ocr


def image_ocr( parse_out_dir, image_saved_dir, ocr):
    """
    函数 image_ocr 读取解析目录中的图像文件信息，执行 OCR 操作并将结果保存。
    参数说明：
    - parse_out_dir: 解析后的输出目录，包含图像信息的 JSON 文件。
    - image_saved_dir: 图像文件保存的目录，用于 OCR 处理。
    - seal_ocr: 是否开启印章检测
    """

    iter_file = traverse_dir(parse_out_dir, ['json'])
    # ofd(发票)和txt(纯文本)文件可以解析，不需要OCR处理
    iter_file = {item for item in iter_file if item[-1] not in {'ofd_content', 'txt_content'}}
    if not iter_file:
        return

    # need to run only once to download and load model into memory
    # 遍历指定目录下的所有 JSON 文件
    exclude_content_set = {'CT)'}
    # 对于每个 JSON 文件，读取内容并执行 OCR 操作
    for json_path, json_name in iter_file:
        if json_name == 'sct_content':
            continue
        # 打开 JSON 文件并加载数据
        with open(json_path, 'r', encoding='utf-8') as file:
            data_dict = json.load(file)

        # 遍历文件列表中的每个图像信息字典
        for dic in data_dict['file_list']:
            image_info = dic['image_info']  # 获取图像信息字典
            uuid = dic['file_uuid']
            if image_info:
                bar = tqdm(image_info.items())
                # 遍历每个图像项，执行 OCR 操作并更新内容
                for k, v in bar:
                    bar.set_description(f'对{uuid[0:8]}进行OCR')
                    img_path = os.path.join(image_saved_dir, v['path'])  # 构建图像完整路径
                    # seal_text = '\n'.join(seal_ocr(img_path)).strip()
                    # if len(seal_text) > 0:
                    #     print('印章文字：', seal_text)
                    try:
                        result = ocr(img_path, use_cls=True)
                        result = [[box.tolist(), [txt, score]] for txt, box, score in
                                  zip(result.txts, result.boxes, result.scores)]
                    except Exception as e:
                        # print(e)
                        result = None
                    if result:
                        # 从图像中提取文本，进行正则过滤并去除空格
                        ocr_text = re_filter(' '.join([line[1][0] for line in result])).strip()
                    else:
                        ocr_text = ''
                    if ocr_text in exclude_content_set:
                        image_info.pop(k)
                    else:
                        image_info[k]['content'] = result  # + '\n印章：\n' + seal_text

                        # 将处理后的数据重新写入 JSON 文件
        with open(json_path, 'w+', encoding='utf-8') as f:
            json.dump(data_dict, f, ensure_ascii=False)
